// Copyright 2024 Andy Fingerhut
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

#ifndef SAI_PACKET_REWRITES_P4_
#define SAI_PACKET_REWRITES_P4_

#include <v1model.p4>
#include "headers.p4"
#include "metadata.p4"
#include "minimum_guaranteed_sizes.p4"
#include "bmv2_intrinsics.h"
#include "ids.h"

// To be applied only for multicast-replicated packets, i.e. packets with
// `standard_metadata.instance_type == PKT_INSTANCE_TYPE_REPLICATION`.
// In P4Runtime, these are packets created by a `Replica` of a
// `MulticastGroupEntry`.
control multicast_rewrites(inout local_metadata_t local_metadata,
                           in standard_metadata_t standard_metadata) {
  // The egress port of the multicast-replicated packet.
  // In P4Runtime, equal to the `port` value of the `Replica` of the
  // `MulticastGroupEntry` that created this packet.
  port_id_t multicast_replica_port = (port_id_t) standard_metadata.egress_port;

  // The instance number of the multicast-replicated packet.
  // In P4Runtime, equal to the `instance` value of the `Replica` of the
  // `MulticastGroupEntry` that created this packet.
  replica_instance_t multicast_replica_instance =
      standard_metadata.egress_rid;

  @id(ROUTING_SET_MULTICAST_SRC_MAC_ACTION_ID)
  action set_multicast_src_mac(@id(1) @format(MAC_ADDRESS)
                               ethernet_addr_t src_mac) {
    local_metadata.enable_src_mac_rewrite = true;
    local_metadata.packet_rewrites.src_mac = src_mac;
    // Hard-coded for now. We may make the VLAN ID programmable when needed.
    local_metadata.packet_rewrites.vlan_id = INTERNAL_VLAN_ID;
  }

  @unsupported
  @id(ROUTING_L2_MULTICAST_PASSTHROUGH_ACTION_ID)
  action l2_multicast_passthrough() {}

  // This is a logical table that does not exist in SAI and instead is managed
  // by the Orchagent. It is used to distinguish between L2 and IP
  // multicast-replicated packets.
  //  * L2MC packets will use the l2_multicast_passthrough action
  //  * IPMC packets will use set_multicast_src_mac action.
  //
  // L2MC packets will not be modified in any way and simply duplicated to
  // various output ports. IP packets will rewrite the source MAC.
  //
  // For IPMC there is a many-to-one correspondence between entries in this
  // table and SAI Router Interfaces (RIFs). Each entry corresponds to a RIF
  // with the following attributes:
  // * `SAI_ROUTER_INTERFACE_ATTR_PORT_ID` is equal to `multicast_replica_port`.
  // * `SAI_ROUTER_INTERFACE_ATTR_SRC_MAC_ADDRESS` is equal to the `src_mac`
  //   parameter of the `set_multicast_src_mac` action.
  //
  // Orchagent maintains a mapping from entries to RIFs, creating and destroying
  // (possibly shared) RIFs as entries are inserted and deleted.
  //
  // When creating a multicast group member (`SAI_IPMC_GROUP_MEMBER`) from a
  // P4Runtime `Replica`, Orchagent will use this table to set the value of the
  // `SAI_IPMC_GROUP_MEMBER_ATTR_IPMC_OUTPUT_ID` attribute: it will expect to
  // find an entry for the replica's port and instance in this table, and will
  // use the ID of the RIF associated with that entry. This will cause the
  // source MAC of packets generated by the group member to be rewritten to the
  // `src_mac` of the `set_multicast_src_mac` action of the entry.
  @p4runtime_role(P4RUNTIME_ROLE_ROUTING)
  @id(ROUTING_MULTICAST_ROUTER_INTERFACE_TABLE_ID)
  table multicast_router_interface_table {
    key = {
      multicast_replica_port : exact
        @referenced_by(builtin::multicast_group_table, replica.port)
        @id(1);
      multicast_replica_instance : exact
        @referenced_by(builtin::multicast_group_table, replica.instance)
        @id(2);
    }
    actions = {
      @proto_id(1) set_multicast_src_mac;
      @proto_id(2) l2_multicast_passthrough;
    }
    size = ROUTING_MULTICAST_SOURCE_MAC_TABLE_MINIMUM_GUARANTEED_SIZE;
  }

  apply {
    multicast_router_interface_table.apply();
  }
}  // control multicast_rewrites

// This control block applies the rewrites computed during the ingress
// stage to the actual packet.
control packet_rewrites(inout headers_t headers,
                        inout local_metadata_t local_metadata,
                        inout standard_metadata_t standard_metadata) {
  apply {
    if (standard_metadata.instance_type == PKT_INSTANCE_TYPE_REPLICATION) {
      local_metadata.enable_decrement_ttl = true;
      multicast_rewrites.apply(local_metadata, standard_metadata);
    }
    if (local_metadata.enable_src_mac_rewrite) {
      headers.ethernet.src_addr = local_metadata.packet_rewrites.src_mac;
    }
    if (local_metadata.enable_dst_mac_rewrite) {
      headers.ethernet.dst_addr = local_metadata.packet_rewrites.dst_mac;
    }
    if (local_metadata.enable_vlan_rewrite) {
      // VLAN id is kept in local_metadata until the end of egress pipeline
      // where depending on the value of VLAN id and VLAN configuration the
      // packet might potentially get VLAN tagged with that VLAN id.
      local_metadata.vlan_id = local_metadata.packet_rewrites.vlan_id;
    }
    if (headers.ipv4.isValid()) {
      if (headers.ipv4.ttl > 0 && local_metadata.enable_decrement_ttl) {
        headers.ipv4.ttl = headers.ipv4.ttl - 1;
      }
      // TODO: Verify this is accurate when TTL rewrite is
      // disabled and update this code if not.
      if (headers.ipv4.ttl == 0) mark_to_drop(standard_metadata);
    }
    if (headers.ipv6.isValid()) {
      if (headers.ipv6.hop_limit > 0 && local_metadata.enable_decrement_ttl) {
        headers.ipv6.hop_limit = headers.ipv6.hop_limit - 1;
      }
      // TODO: Verify this is accurate when TTL rewrite is
      // disabled and update this code if not.
      if (headers.ipv6.hop_limit == 0) mark_to_drop(standard_metadata);
    }
  }
}  // control packet_rewrites

#endif  // SAI_PACKET_REWRITES_P4_
